import pandas as pd
import numpy as np
import torch

from sklearn.metrics import mean_squared_error
from tqdm import tqdm

# 读取数据集
train_data = pd.read_csv('./dataset/train_data.csv')
similarity_matrix_save_path = 'E:/run_res/similarity_matrix_shuffle.csv'
similarity_matrix = pd.read_csv(similarity_matrix_save_path, index_col=0)

df = pd.DataFrame(train_data)
# 物品-用户评分矩阵
user_item_matrix = pd.pivot_table(df, values='rating', index='user_id', columns='book_id', fill_value=0)

def predict_all_user_ratings(user_item_matrix, similarity_matrix):
    # 获取用户对物品的评分
    similar_items = torch.tensor(similarity_matrix.values, dtype=torch.float32).cuda()  # 将张量移动到 GPU 上

    all_predicted_ratings = []

    # 为每个用户生成预测评分
    for user_id in tqdm(user_item_matrix.index, desc="Generating Predictions", unit="user"):
        user_ratings = torch.tensor(user_item_matrix.loc[user_id].values, dtype=torch.float32).cuda()  # 将张量移动到 GPU 上
 
        # 计算用户已评分项的相似度之和
        sum_of_similarities = torch.abs(similar_items[:, user_ratings != 0]).sum(dim=1)

        # 修改分母，仅考虑用户已评分项的相似度之和
        predicted_ratings = torch.mv(similar_items.T, user_ratings) / sum_of_similarities

        all_predicted_ratings.append(predicted_ratings.cpu().numpy())  # 将结果移回 CPU

    return np.array(all_predicted_ratings)

# 为所有用户生成预测评分矩阵
all_user_predictions = predict_all_user_ratings(user_item_matrix, similarity_matrix)

# 保存预测评分矩阵到文件
# np.savetxt('E:/run_res/all_user_ratings_predictions_pytorch_gpu.csv', all_user_predictions, delimiter=',')
# print("所有用户的预测评分已保存。")
test_data = pd.read_csv('./dataset/test_data.csv')
predicted_ratings_matrix = np.vstack(all_user_predictions)

# 提取测试集中的用户、物品和真实评分
user_ids_test = test_data['user_id'].values
book_ids_test = test_data['book_id'].values
ratings_test = test_data['rating'].values

# 通过用户和物品的索引获取预测评分
predicted_ratings_test = predicted_ratings_matrix[user_ids_test - 1, book_ids_test - 1]

# 计算均方误差（MSE）
mse = mean_squared_error(ratings_test, predicted_ratings_test)


print(f'MSE: {mse}')